# load standard libraries #
import random
import time
from math import *

# load C-side code #
from OpenNero import *

# load Python functional scripts #
from common import *
import world_handler
from agent_handler import AgentState, AgentInit

# load agent script
from roomba import RoombaBrain

from constants import *

# imports for profile project
from scipy.cluster.vq import kmeans
from numpy import array
from pybrain.optimization import GA

import copy
import pickle
import pprint
import os.path
# end imports for profile project

class SandboxMod:

    def __init__(self):
        """
        initialize the sandbox server
        """
        self.marker_map = {} # a map of cells and markers so that we don't have more than one per cell
        self.environment = None
        self.agent_ids = []
    
    def mark(self, x, y, marker):
        """ Mark a position (x, y) with the specified color """
        # remove the previous object, if necessary
        self.unmark(x, y)
        # add a new marker object
        id = addObject(marker, Vector3f(x, y, -1), Vector3f(0,0,0), Vector3f(0.5,0.5,0.5), type = OBJECT_TYPE_MARKER)
        # remember the ID of the object we are about to create
        self.marker_map[(x, y)] = id
	    
    def mark_blue(self, x, y):
        self.mark(x, y,"data/shapes/cube/BlueCube.xml")
    
    def mark_green(self, x, y):
        self.mark(x, y,"data/shapes/cube/GreenCube.xml")
    
    def mark_yellow(self, x, y):
        self.mark(x, y,"data/shapes/cube/YellowCube.xml")
    
    def mark_white(self, x, y):
        self.mark(x, y,"data/shapes/cube/WhiteCube.xml")
    
    def unmark(self, x, y):
        if (x, y) in self.marker_map:
            removeObject(self.marker_map[(x, y)])
            del self.marker_map[(x, y)]
            return True
        else:
            return False
    
    def setup_sandbox(self):
        """
        setup the sandbox environment
        """
        global XDIM, YDIM, HEIGHT, OFFSET
        self.environment = RoombaEnvironment(XDIM, YDIM)
        set_environment(self.environment)
    
    def reset_sandbox(self=None):
        """
        reset the sandbox and refill with stuff to vacuum
        """
        for id in self.marker_map.values():
            removeObject(id)  # delete id from Registry, not from dict
        self.marker_map = {}
        for id in self.agent_ids:
            removeObject(id)  # delete id from Registry, not from list
        self.agent_ids = []
        reset_ai()

    def remove_bots(self):
        """ remove all existing bots from the environment """
        disable_ai()
        for id in self.agent_ids:
            removeObject(id)  # delete id from Registry, not from list
        self.agent_ids = []

    def distribute_bots(self, num_bots, bot_type):
        """distribute bots so that they don't overlap"""
        # make a number of tiles to stick bots in
        N_TILES = 10
        tiles = [ (r,c) for r in range(N_TILES) for c in range(N_TILES)]
        if not self.environment.SEED is None:
            random.seed(self.environment.SEED)
        random.shuffle(tiles)
        for i in range(num_bots):
            (r,c) = tiles.pop() # random tile
            x, y = r * XDIM / float(N_TILES), c * YDIM / float(N_TILES) # position within tile
            if not self.environment.SEED is None:
                random.seed(self.environment.SEED+1000)
                x_rand = random.random()
                random.seed(self.environment.SEED-1000)
                y_rand = random.random()
            else:
                x_rand = random.random()
                y_rand = random.random()
            x, y = x + x_rand * XDIM * 0.5 / N_TILES, y + y_rand * YDIM * 0.5 / N_TILES # random offset
            agent_id = addObject(bot_type, Vector3f(x, y, 0), Vector3f(0.5, 0.5, 0.5), type = OBJECT_TYPE_ROOMBA, collision = OBJECT_TYPE_ROOMBA)
            self.agent_ids.append(agent_id)
        
    def add_bots(self, bot_type, num_bots):
        disable_ai()
        num_bots = int(num_bots)
        if bot_type.lower().find("script") >= 0:
            self.distribute_bots(num_bots, "data/shapes/roomba/Roomba.xml")
            enable_ai()
            return True
        else:
            return False

def in_bounds(x,y):
    return x > ROOMBA_RAD and x < XDIM - ROOMBA_RAD and y > ROOMBA_RAD and y < YDIM - ROOMBA_RAD

#################################################################################        
class RoombaEnvironment(Environment):
    """
    Sample Environment for the Sandbox
    """
    def __init__(self, XDIM, YDIM):
        """
        Create the environment
        """
        Environment.__init__(self)

        # Profile project additions
        # Constants
        self.VERBOSE = False
        self.NUM_OF_CRUMBS = 500
        self.NUM_OF_AGENTS = 5
        self.NUM_OF_CLUSTERS = 10
        self.GRID_X = 2
        self.GRID_Y = 5
        self.MAX_EVAL = 250000
        self.LOOKAHEAD = False
        self.EMPTY_ROOM_CONDITION = True
        self.SEED = 0
        
        # Vars
        self.centroids = None
        self.centroid_assignment = None
        self.log_ep = -1
        self.heuristic_indices = []
        self.min_dist = float("inf")
        self.bumps = 0
        self.cluster_list = None
        self.time1 = 0
        self.time2 = 0
        
        # 0 = grid clustering
        # 1 = kMeans clustering
        self.CLUSTER_MODE = 0
        # 0 = brute force assignment
        # 1 = Genetic algorithm assignment
        # 2 = Heuristic assignment
        # 3 = GA future look
        self.ASS_MODE = 1
        # 0 = distance only
        # 1 = pellets only
        # 2 = heuristic combination
        self.GA_eval = 2
        # 0 = norm(pellets) - norm(dist)
        # 1 = max_dist-dist
        # 2 = pellets
        self.HEURISTIC = 0
        
        if self.ASS_MODE == 3:
            self.MAX_EVAL = 2500000
        
        ### Create logfile name
        third_number = 0
        if(self.ASS_MODE == 2):
            third_number = self.HEURISTIC
        else:
            third_number = self.GA_eval
            
        if(self.CLUSTER_MODE == 0):
            self.LOGNAME = 'log_' + str(self.CLUSTER_MODE) + str(self.ASS_MODE) + str(third_number) + '(' + str(self.GRID_X) + 'x' + str(self.GRID_Y) + ')_mutstd3.txt'
        if(self.CLUSTER_MODE == 1):
            self.LOGNAME = 'log_' + str(self.CLUSTER_MODE) + str(self.ASS_MODE) + str(third_number) + '(' + str(self.NUM_OF_CLUSTERS) + ').txt'
        ###
        
        try:
            file1 = open('list' + str(self.CLUSTER_MODE) + str(self.ASS_MODE) + str(third_number) + '.pck', 'r')
            list = pickle.load(file1)
            file1.close()
            self.SEED = list.pop()
            file1 = open('list' + str(self.CLUSTER_MODE) + str(self.ASS_MODE) + str(third_number) + '.pck', 'w')
            pickle.dump(list, file1)
            file1.close()
        except:
            file1 = open('list' + str(self.CLUSTER_MODE) + str(self.ASS_MODE) + str(third_number) + '.pck', 'w')
            list = [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29]
            pickle.dump(list, file1)
            file1.close()
            self.SEED = 0
        print 'SEED:', self.SEED
        
        if self.CLUSTER_MODE == 0:
            self.NUM_OF_CLUSTERS = self.GRID_X * self.GRID_Y
        if self.ASS_MODE == 3:
            self.LOOKAHEAD = True
        # End profile project additions
        
        self.XDIM = XDIM
        self.YDIM = YDIM
        self.max_steps = 10000

        print 'sanity check'
        
        self.states = {} # dictionary of agent states
        self.crumbs = world_handler.pattern_cluster(500, self.SEED) #"Roomba/world_config.txt"
        # only keep crumbs that are inside the walls
        #self.crumbs = [c for c in self.crumbs if in_bounds(c.x,c.y)]
        # self.crumbs = world_handler.fill_random_pellet(self.NUM_OF_CRUMBS, XDIM, YDIM, self.SEED)
        
        print 'sanity check 2'
        
        
        self.init_list = AgentInit()
        self.init_list.add_type("<class 'Roomba.roomba.RoombaBrain'>")
        #print self.init_list.types

        roomba_abound = self.init_list.get_action("<class 'Roomba.roomba.RoombaBrain'>")
        roomba_sbound = self.init_list.get_sensor("<class 'Roomba.roomba.RoombaBrain'>")
        roomba_rbound = self.init_list.get_reward("<class 'Roomba.roomba.RoombaBrain'>")

        ### Bounds for Roomba ###
        # actions
        roomba_abound.add_continuous(-pi, pi) # amount to turn by
        
        # sensors
        roomba_sbound.add_discrete(0,1)    # wall bump
        roomba_sbound.add_continuous(0,XDIM)   # self.x
        roomba_sbound.add_continuous(0,YDIM)   # self.y
        roomba_sbound.add_continuous(0,XDIM)   # closest.x
        roomba_sbound.add_continuous(0,YDIM)   # closest.y
        
        # rewards
        roomba_rbound.add_continuous(-100,100) # range for reward

        ### End Bounds for Roomba ####

        # set up shop
        # Add Wayne's Roomba room with experimentally-derived vertical offset to match crumbs.
        addObject("data/terrain/RoombaRoom.xml", Vector3f(XDIM/2,YDIM/2, -1), Vector3f(0,0,0), Vector3f(XDIM/245.0, YDIM/245.0, HEIGHT/24.5), type = OBJECT_TYPE_WALLS)

        # getSimContext().addAxes()
        self.add_crumbs()
        for crumb in self.crumbs:
            self.add_crumb_sensors(roomba_sbound)       

        # log time
        self.time1 = time.clock()
        
        print 'INIT COMPLETE'

    def get_state(self, agent):
        if agent in self.states:
            return self.states[agent]
        else:
            # print "new state created"
            pos = agent.state.position
            rot = agent.state.rotation
            self.states[agent] = AgentState(pos, rot)
            return self.states[agent]
        
    #def randomize(self):
    #    self.crumbs = world_handler.read_pellets()
    #    # only keep crumbs that are inside the walls        
    #    self.crumbs = [c for c in self.crumbs if in_bounds(c.x,c.y)]

    def add_crumb_sensors(self, roomba_sbound):
        """Add the crumb sensors, in order: x position of crumb (0 to XDIM,
        continuous), y position of crumb (0 to XDIM, continuous), whether
        crumb is present at said position or has been vacced (0 or 1), and
        reward for vaccing said crumb."""
        roomba_sbound.add_continuous(0, XDIM)    # crumb position X
        roomba_sbound.add_continuous(0, YDIM)    # crumb position Y
        roomba_sbound.add_discrete(0, 1)       # crumb present/ not present
        roomba_sbound.add_discrete(1, 5)       # reward for crumb

    def add_crumbs(self):
        for pellet in self.crumbs:
            if not (pellet.x, pellet.y) in getMod().marker_map:
                getMod().mark_blue(pellet.x, pellet.y)

    def reset(self, agent):
        """ reset the environment to its initial state """
        state = self.get_state(agent)
        if self.VERBOSE:
            print "Episode %d complete" % state.episode_count
        
        # Write performance to log
        self.time2 = time.clock()
        remaining_crumbs = 0
        for c in self.crumbs:
            if (c.x, c.y) in getMod().marker_map:
                remaining_crumbs += 1
                getMod().unmark(c.x, c.y)
        collected = len(self.crumbs) - remaining_crumbs
        steps = state.step_count
        time3 = self.time2 - self.time1
        to_write = str(collected) + '   ,   ' + str(steps) + '   ,   ' + str(time3) + '\n'
        print "Collected crumbs: ", collected
        print "Steps: ", steps
        f = open('Roomba/' + self.LOGNAME, 'a')
        f.write(to_write)
        f.close()
        
        getMod().remove_bots()
        getMod().add_bots('script', self.NUM_OF_AGENTS)
        self.__init__(self.XDIM, self.YDIM)

    def get_agent_info(self, agent):
        """ return a blueprint for a new agent """
        return self.init_list.get_info(str(type(agent)))

    def num_sensors(self):
        """
        Return total number of sensors
        """
        return (len(getMod().marker_map)*4 + N_FIXED_SENSORS)
    
    def step(self, agent, action):
        """
        A step for an agent
        """
        state = self.get_state(agent) # the agent's status
        reward = 0
        if (state.is_out == True):
            getMod().unmark(agent.state.position.x, agent.state.position.y)
        else:
            assert(self.get_agent_info(agent).actions.validate(action)) # check if the action is valid
            if (str(type(agent)) == "<class 'Roomba.roomba.RoombaBrain'>"):
                angle = action[0] # in range of -pi to pi
                degree_angle = degrees(angle)
                delta_angle = degree_angle - agent.state.rotation.z
                delta_dist = MAX_SPEED
            else:
                # The first action specifies the distance to move in the forward direction
                # and the second action specifies the angular change in the orientation of
                # the agent.
                delta_dist = action[0]*MAX_SPEED
                delta_angle = action[1]*ANGULAR_SPEED
            if self.detect_bump(agent, delta_dist, delta_angle):
                delta_dist = -5*delta_dist #5 is a good val
                delta_angle = delta_angle + radians(90)
            reward = self.update_position(agent, delta_dist, delta_angle)
        state.reward += reward
        return reward
        
    def detect_bump(self, agent, delta_dist, delta_angle):
        position = agent.state.position
        rotation = agent.state.rotation

        # posteriori collision detection
        rotation.z = wrap_degrees(rotation.z, delta_angle)
        position.x += delta_dist*cos(radians(rotation.z))
        position.y += delta_dist*sin(radians(rotation.z))
        for other_agent in self.states:
            if(other_agent != agent):
                dist = sqrt((other_agent.state.position.x - position.x) ** 2 + (other_agent.state.position.y - position.y) ** 2)
                if(dist <= (2 * ROOMBA_RAD) + 2):
                    if self.VERBOSE:
                        print "Bump Detected!"
                    self.bumps += 1
                    self.update_position(other_agent, -5,90)
                    
                    return True
        return False

    # delta_angle (degrees) is change in angle
    # delta_dist is change in distance (or velocity, since unit of time unchanged)
    def update_position(self, agent, delta_dist, delta_angle):
        """
        Updates position of the agent and collects pellets.
        """
        state = self.get_state(agent)
        state.step_count += 1

        position = agent.state.position
        rotation = agent.state.rotation

        # posteriori collision detection
        rotation.z = wrap_degrees(rotation.z, delta_angle)
        position.x += delta_dist*cos(radians(rotation.z))
        position.y += delta_dist*sin(radians(rotation.z))

        # check if one of 4 out-of-bound conditions applies
        # if yes, back-track to correct position
        if (position.x) < 0 or (position.y) < 0 or \
           (position.x) > self.XDIM or (position.y) > self.YDIM:

            # correct position
            if (position.x) < 0:
                position.x -= delta_dist*cos(radians(rotation.z))    
            if (position.y) < 0:
                position.y -= delta_dist*sin(radians(rotation.z))
            if (position.x) > self.XDIM:
                position.x -= delta_dist*cos(radians(rotation.z))
            if (position.y) > self.YDIM:
                position.y -= delta_dist*sin(radians(rotation.z))
            
        # register new position
        state.position = position
        state.rotation = rotation
        agent.state.position = position
        agent.state.rotation = rotation
        
        reward = 0
        
        # remove all crumbs within ROOMBA_RAD of agent position
        pos = (position.x, position.y)
        for crumb in self.crumbs:
            if (crumb.x, crumb.y) in getMod().marker_map:
                distance = sqrt((crumb[0] - pos[0])**2 + (crumb[1] - pos[1])**2)
                if distance < ROOMBA_RAD:
                    getMod().unmark(crumb.x, crumb.y)
                    reward += 1
                
        # check if agent has expended its step allowance
        if (self.max_steps != 0) and (state.step_count >= self.max_steps):
            state.is_out = True    # if yes, mark it to be removed
        return reward

    def grid_cluster(self):
        # Create a grid of clusters for N agents without kMeans
        if self.VERBOSE:
            print "Clustering using grid"
        x_step = self.XDIM / self.GRID_X
        y_step = self.YDIM / self.GRID_Y
        centroids = []
        for x in range(self.GRID_X):
            for y in range(self.GRID_Y):
                x_coord = x*x_step+x_step/2
                y_coord = y*y_step+y_step/2
                centroids.append((x_coord,y_coord))
        return centroids
        
    def k_means_cluster(self):
        #print "nr of crumbs:", len(self.crumbs)
        if self.VERBOSE:
            print "Clustering with k-means"    
        pellets = []
        for pellet in self.crumbs:
            if (pellet.x, pellet.y) in getMod().marker_map:
                pellets.append([pellet.x, pellet.y])
        pelletArray = array(pellets)
        [centroids, distortion] = kmeans(pelletArray, self.NUM_OF_CLUSTERS)
        while len(centroids) != self.NUM_OF_CLUSTERS:
            print "centroid size = ", len(centroids), ". Reclustering..."
            [centroids, distortion] = kmeans(pelletArray, self.NUM_OF_CLUSTERS)
        return centroids
        
    def list_clusters(self):
        if self.VERBOSE:
            print "Creating separate cluster pellet lists..."
        cluster_list = []
        for i in range(len(self.centroids)):
            cluster_list.append([])
        for pellet in self.crumbs:
            min_dist = float('inf')
            min_pellet = None
            x = pellet.x
            y = pellet.y
            if (x,y) in getMod().marker_map:
                for i in range(len(self.centroids)):
                    dist = self.calc_euclidean([x,y], self.centroids[i])
                    if dist < min_dist:
                        min_dist = dist
                        min_pellet = i
                cluster_list[min_pellet].append((x,y))
        return cluster_list
    
    def sense(self, agent, sensors):
        # sensors[0] is not used
        sensors[0] = 0
        
        # get agent's position
        pos = agent.state.position
        sensors[1] = pos.x
        sensors[2] = pos.y

        # only start sensing if all agents have been initialized
        if len(self.states) == self.NUM_OF_AGENTS:
            # cluster if this has not yet been done
            if self.centroids == None:
                # grid clustering or k-means clustering
                if self.CLUSTER_MODE == 0:
                    self.centroids = self.grid_cluster()
                elif self.CLUSTER_MODE == 1:
                    self.centroids = self.k_means_cluster()
                    
                # draw centroids
                for centroid in self.centroids:
                    getMod().mark_green(centroid[0], centroid[1])
                    
                # create lists of pellets belonging to centroids
                self.cluster_list = self.list_clusters()
                
                self.reassign(agent)
            
            agent_centroid = self.get_centroid(agent)
            #print 'agent', agent.state.position.x, agent.state.position.y
            #print 'assignment', agent_centroid
            #print self.centroids
            if self.LOOKAHEAD:
                ass = None
                for assignment in agent_centroid:
                    if len(self.cluster_list[assignment]) > 0:
                        if ass == None:
                            ass = assignment
                if ass == None:
                    ass = agent_centroid[len(agent_centroid)-1]
                self.sense_crumbs(sensors, N_S_IN_BLOCK, N_FIXED_SENSORS, agent, self.cluster_list[ass])
            else:
                if len(self.cluster_list[agent_centroid]) == 0:
                    non_empty_clusters = self.get_non_empty_clusters()
                    for c in non_empty_clusters:
                        if not c in self.centroid_assignment:
                            if self.VERBOSE:
                                print "Cluster", agent_centroid, "is empty, reassigning"
                            self.reassign(agent)
                            agent_centroid = self.get_centroid(agent)
                            if len(self.cluster_list[agent_centroid]) == 0:
                                agent_centroid = self.assign_closest_cluster(non_empty_clusters, agent, agent_centroid)
                            self.heuristic_indices = []
                            break
                self.sense_crumbs(sensors, N_S_IN_BLOCK, N_FIXED_SENSORS, agent, self.cluster_list[agent_centroid])
        return sensors
        
    def assign_closest_cluster(self, list, agent, agent_centroid):
        min_dist = float('inf')
        state = self.get_state(agent)
        x = state.position.x
        y = state.position.y
        cluster = None
        for i in list:
            if not i in self.centroid_assignment:
                centroid = self.centroids[i]
                dist = self.calc_euclidean([x,y], centroid)
                if dist < min_dist:
                    min_dist = dist
                    cluster = i
        new_assignments = []
        for ass in self.centroid_assignment:
            if ass[1] == agent_centroid:
                new_assignments.append((0, cluster, [x,y]))
            else:
                new_assignments.append(ass)
        self.centroid_assignment = new_assignments
        return cluster
    
    def reassign(self, agent):
        if self.ASS_MODE == 0:
            # self.centroid_assignment = 
            self.brute_force_assignment()
        elif self.ASS_MODE == 1:
            centroid_assignment = self.GA_assignment()
            #while self.contains_out_of_bound_elements(self.centroid_assignment, 0, len(self.centroids) - 1) or self.hasDuplicates(self.centroid_assignment):
            #    self.centroid_assignment = self.GA_assignment()
        elif self.ASS_MODE == 2:
            centroid_assignment = self.heuristic_assignment()
        elif self.ASS_MODE == 3:
            centroid_assignment = self.GA_future()
        
        if self.centroid_assignment != None and not self.LOOKAHEAD:
            self.centroid_assignment = self.replace_assignment(centroid_assignment, agent)
        else:
            self.centroid_assignment = centroid_assignment
    
    def replace_assignment(self, assignments, agent):
        state = self.get_state(agent)
        x = state.position.x
        y = state.position.y
        toAdd = None
        closest_dist = float('inf')
        for ass in assignments:
            dist = self.calc_euclidean([x, y], ass[2])
            if dist <  closest_dist:
                closest_dist = dist
                toAdd = ass
        if toAdd == None:
            print assignments
            print x,y
        new_assignments = []
        agent_centroid = self.get_centroid(agent)
        for ass in self.centroid_assignment:
            if ass[1] == agent_centroid:
                new_assignments.append(toAdd)
            else:
                new_assignments.append(ass)
        return new_assignments
    
    def get_non_empty_clusters(self):
        list =[]
        for cluster in self.cluster_list:
            if len(cluster) > 0:
                list.append(self.cluster_list.index(cluster))
        return list
    
    def calc_max_pellets(self):
        self.max_pellets = 0
        for cluster in self.cluster_list:
            pellets = len(cluster)
            if pellets > self.max_pellets:
                self.max_pellets = pellets
    
    def brute_force_assignment(self):
        if self.VERBOSE:
            print "Assignment by brute-force"
            # calculate number of combinations to calculate for brute-force assignment
            num_centroids = len(self.centroids)
            prod = 1
            for i in range(5):
                prod = prod * num_centroids
                num_centroids = num_centroids - 1
            print 'Number of combinations to calculate:', prod
        agents = self.get_all_agents()
        distances = self.create_distance_matrix()
        agent_list = range(len(agents))
        centroid_list = range(len(self.centroids))
        self.find_shortest_path(distances, agent_list, centroid_list, [], agents)
        self.min_dist = float("inf")
    
    def find_shortest_path(self, matrix, agent_list, centroid_list, assignments, agents):
        if len(assignments) == len(self.states):
            dist = self.calc_distance(matrix, assignments)
            if dist < self.min_dist:
                self.min_dist = dist
                self.centroid_assignment = assignments
        else:
            for i in agent_list:
                sub_agent_list = copy.deepcopy(agent_list)
                sub_agent_list.remove(i)
                for j in centroid_list:
                    sub_centroid_list = copy.deepcopy(centroid_list) 
                    sub_centroid_list.remove(j)
                    ass_copy = copy.deepcopy(assignments)
                    ass_copy.append((i,j,agents[i]))
                    self.find_shortest_path(matrix, sub_agent_list, sub_centroid_list, ass_copy, agents)
    
    def GA_future(self):
        if self.VERBOSE:
            "Assignment by genetic algorithm, with lookahead"
        
        clusters_per_agent = len(self.centroids) / len(self.states)
        self.distances = self.create_distance_matrix()
        self.calc_max_pellets()
        
        # initialize GA with random assignment
        ass_list = range(len(self.centroids))
        x0 = random.sample(ass_list, len(self.centroids))
        
        # initialize genetic algorithm and run it
        gen_alg = GA(self.eval_future, x0)
        gen_alg.maxEvaluations = self.MAX_EVAL
        #gen_alg.minimize = True
        assignments = gen_alg.learn()
        
        agents = self.get_all_agents()
        final_assignments = []
        for i in range(len(self.states)):
            sub_assignment = []
            for j in range(clusters_per_agent):
                sub_assignment.append(int(assignments[0][i*clusters_per_agent+j]))
            final_assignments.append((i, sub_assignment, agents[i]))
        return final_assignments
    
    def GA_assignment(self):
        if self.VERBOSE:
            print "Assignment by genetic algorithm"
        
        self.distances = self.create_distance_matrix()
        self.calc_max_pellets()
                
        # initialize GA with random assignment
        ass_list = range(self.NUM_OF_CLUSTERS)
        x0 = random.sample(ass_list,self.NUM_OF_AGENTS)
                    
        # initialize genetic algorithm and run it
        gen_alg = GA(self.eval, x0)
        gen_alg.maxEvaluations = self.MAX_EVAL
        assignments = gen_alg.learn()
        
        int_list = []
        for a in assignments[0]:
            int_list.append(int(a))
        
        while self.contains_out_of_bound_elements(int_list,0,self.NUM_OF_CLUSTERS-1) :
            print "Out of bounds assignment!  val: ", assignments[1]
            print "assignments:", int_list
            assignments = gen_alg.learn()
            int_list = []
            for a in assignments[0]:
                int_list.append(int(a))
                
                
        agents = self.get_all_agents()
        
        final_assignments = []
        for i in range(len(self.states)):
            final_assignments.append((i, int_list[i], agents[i]))
        return final_assignments
    
    # Makes a matrix of rows * columns
    # Input: number of rows, number of columns
    # Returns: a two dimensional list usable as matrix
    def make_matrix(self, rows, cols):
        matrix = []
        for i in range(rows):
            matrix.append([])
            for j in range(cols):
                matrix[i].append([])
        return matrix
        
    # create a distance matrix for every agent to every centroid
    def create_distance_matrix(self):
        if self.VERBOSE:
            print 'Creating distance matrix...'
        agents = self.get_all_agents()
        distances = self.make_matrix(len(agents), len(self.centroids))
        for i in range(len(agents)):
            agent_x = agents[i][0]
            agent_y = agents[i][1]
            for j in range(len(self.centroids)):
                centroid_x = self.centroids[j][0]
                centroid_y = self.centroids[j][1]
                dist = sqrt(pow(agent_x - centroid_x, 2) + pow(agent_y - centroid_y, 2))
                distances[i][j] = dist
        return distances
        
    # create a distance matrix for every agent to every centroid
    def create_distance_matrix(self):
        if self.VERBOSE:
            print 'Creating distance matrix...'
        self.max_distance = 0
        agents = self.get_all_agents()
        distances = self.make_matrix(len(agents), len(self.centroids))
        for i in range(len(agents)):
            agent_x = agents[i][0]
            agent_y = agents[i][1]
            for j in range(len(self.centroids)):
                centroid_x = self.centroids[j][0]
                centroid_y = self.centroids[j][1]
                dist = sqrt(pow(agent_x - centroid_x, 2) + pow(agent_y - centroid_y, 2))
                if dist > self.max_distance:
                    self.max_distance = dist
                distances[i][j] = dist
        return distances
        
    def heuristic_assignment(self):
        if self.VERBOSE:
            print "Assignment by heuristics"
        self.centroid_assignment = []
        agents = self.get_all_agents()
        self.calc_max_pellets()
        max_distance = sqrt(self.XDIM ** 2 + self.YDIM ** 2)
        for i in range(len(agents)):
            agent_x = agents[i][0]
            agent_y = agents[i][1]
            iterator = 0
            best_value = float('-inf')
            final_ass = -1
            value = 0
            for centroid in self.centroids:
                if iterator not in self.heuristic_indices:
                    dist = self.calc_euclidean([agent_x, agent_y], [centroid[0], centroid[1]])
                    num_pellets = len(self.cluster_list[iterator])
                    if num_pellets == 0:
                        iterator += 1
                        continue
                    if(self.HEURISTIC == 0):
                    # Heuristic #0: pellets in cluster - dist from agent to cluster
                        norm_dist = dist / max_distance
                        norm_pellets = float(num_pellets) / self.max_pellets
                        if norm_dist != 0:
                            value = norm_pellets / norm_dist
                        else:
                            value = norm_pellets
                    # Heuristic #1: dist from agent to cluster
                    if(self.HEURISTIC == 1):
                        value = max_distance - dist
                    # Heuristic #2: number of pellets in cluster
                    if(self.HEURISTIC == 2):
                        value = num_pellets
                    # print 'cluster', centroid, ' has value', value
                    if value > best_value:
                        best_value = value
                        final_ass = iterator
                iterator = iterator + 1
            self.heuristic_indices.append(final_ass)
            self.centroid_assignment.append((i, final_ass, agents[i]))
        return self.centroid_assignment

    
    def GA_pellets(self, assignments):
        value = 0
        for assignment in assignments:
            centroid = self.cluster_list[assignment]
            value += len(centroid)
        return value           
    
    # define evaluator
    def eval(self, list):
        multiplier = 1
        # convert to ints, due to the discreteness of our problem
        int_list = []
        for i in list:
            int_list.append(int(i))
        # list may not contain duplicate assignments
        if self.hasDuplicates(int_list):
            multiplier = 0.5
        # list may not contain assignments below 0 or above the number of clusters
        # max bound is number of clusters - 1 due to python list assignment goes 
        # from 0 -> length of clusters - 1
        elif self.contains_out_of_bound_elements(int_list, 0, self.NUM_OF_CLUSTERS-1):
            return 0
        # else, return the heuristic reward of the GA
        elif self.GA_eval == 0:
            return self.GA_distance(int_list)*multiplier
        elif self.GA_eval == 1:
            return self.GA_pellets(int_list)*multiplier
        elif self.GA_eval == 2:
            return self.GA_heuristic(int_list)*multiplier
            
    def eval_future(self, list):
        multiplier = 1
        # convert to ints, due to the discreteness of our problem
        int_list = []
        for i in list:
            int_list.append(int(i))
        # list may not contain duplicate assignments
        if self.hasDuplicates(int_list):
            multiplier = 0.5
        # list may not contain assignments below 0 or above the number of clusters
        # max bound is number of clusters - 1 due to python list assignment goes 
        # from 0 -> length of clusters - 1
        if self.contains_out_of_bound_elements(int_list, 0, len(self.centroids)-1):
            return 0
        if self.GA_eval == 0:
            return self.GA_future_distance(int_list)*multiplier
        if self.GA_eval == 1:
            return self.GA_future_pellets(int_list)*multiplier
        if self.GA_eval == 2:
            return self.GA_future_heuristic(int_list)*multiplier
            
    def GA_pellets(self, list):
        value = 0
        for i in list:
            value += len(self.cluster_list[i])
        return value
            
    def GA_distance(self, list):
        value = 0
        max_dist = sqrt(self.XDIM ** 2 + self.YDIM ** 2)
        for i in list:
            dist = self.distances[list.index(i)][i]
            if len(self.cluster_list[i]) > 0:
                value += (max_dist - dist)
        return value

    def contains_out_of_bound_elements(self, list, min, max):
        for i in list:
            if i < min or i > max:
                return True
        return False
    
    def hasDuplicates(self,list):
        set_list = set(list)
        if len(list) == len(set_list):
            return False
        else:
            return True
    
    def GA_heuristic(self, list):
        reward = 0
        for i in list:
            dist = self.distances[list.index(i)][i]
            num_pellets = len(self.cluster_list[i])
            max_distance = sqrt(self.XDIM ** 2 + self.YDIM ** 2)
            norm_dist = dist / max_distance
            norm_pellets = float(num_pellets) / self.max_pellets
            if norm_dist != 0:
                val = norm_pellets / norm_dist
            else:
                val = norm_pellets
            reward = reward + val
        return reward
        
    def GA_future_distance(self, list):
        clusters_per_agent = len(self.centroids) / len(self.states)
        agents = self.get_all_agents()
        # dist
        total_dist = 0
        for agent in agents:
            new_list = []
            for j in range(clusters_per_agent):
                new_list.append(list[agents.index(agent)*clusters_per_agent+j])
            route_dist = self.calc_route_distance(new_list, agent)
            total_dist += route_dist
        return total_dist
        
    def GA_future_heuristic(self, list):
        agents = self.get_all_agents()
        clusters_per_agent = self.NUM_OF_CLUSTERS / self.NUM_OF_AGENTS
        avg_pellets = self.NUM_OF_CRUMBS / self.NUM_OF_AGENTS
        max_dist = sqrt(self.XDIM ** 2 + self.YDIM ** 2) * clusters_per_agent
        max_dist_from_avg = (self.NUM_OF_AGENTS-1)*avg_pellets+(self.NUM_OF_CRUMBS-avg_pellets)
        value = 0
        dist_from_avg = 0
        for agent in agents:
            cluster_list = []
            for j in range(clusters_per_agent):
                cluster_list.append(list[agents.index(agent)*clusters_per_agent+j])
            route_dist = self.calc_route_distance(cluster_list, agent)
            norm_route_dist = route_dist / max_dist
            num_pellets = 0
            for i in cluster_list:
                num_pellets += len(self.cluster_list[i])
            dist_from_avg += abs(avg_pellets - num_pellets)
            norm_dist_from_avg = dist_from_avg / float(max_dist_from_avg)
            value += norm_route_dist + norm_dist_from_avg
        return value
            
    def GA_future_pellets(self, list):
        agents = self.get_all_agents()
        avg_pellets = self.NUM_OF_CRUMBS / self.NUM_OF_AGENTS
        clusters_per_agent = self.NUM_OF_CLUSTERS / self.NUM_OF_AGENTS
        # pellets
        dist_from_avg = 0
        for agent in agents:
            cluster_list = []
            for j in range(clusters_per_agent):
                cluster_list.append(list[agents.index(agent)*clusters_per_agent+j])
            num_pellets = 0
            for i in cluster_list:
                num_pellets += len(self.cluster_list[i])
            dist_from_avg += abs(avg_pellets - num_pellets)
        return 1/float(dist_from_avg)
    
    def calc_route_distance(self, list, agent):
        max_dist = sqrt(self.XDIM ** 2 + self.YDIM ** 2)
        loc_list = [(agent[0], agent[1])]
        for i in list:
            loc_list.append(tuple(self.centroids[i]))
        dist = 0
        for i in loc_list:
            if (loc_list.index(i) < len(loc_list)-1):
                dist += max_dist - self.calc_euclidean(i, loc_list[loc_list.index(i)+1])
        return dist
        
    def get_centroid(self, agent):
        min_dist = float('inf')
        min_ass = None
        for ass in self.centroid_assignment:
            dist = self.calc_euclidean([agent.state.position.x, agent.state.position.y], ass[2])
            if dist < min_dist:
                min_dist = dist
                min_ass = ass
        self.centroid_assignment.remove(min_ass)
        self.centroid_assignment.append([min_ass[0], min_ass[1], [agent.state.position.x, agent.state.position.y]])
        return min_ass[1]
        
    # Calculates the total distance from the distances of the agents to the centroids in assignments.
    # Input: matrix of all distances, assignments of agents to centroids
    # Returns: the total distance    
    def calc_distance(self, matrix, assignments):
        dist = 0
        for a in assignments:
            dist = dist+matrix[a[0]][a[1]]
        return dist
        
    # Finds the location of all agents in the state.
    # Returns: the location of each agent as [[agent1.x, agent1.y], [agent2.x, agent2.y], [..]]    
    def get_all_agents(self):
        all_agents = []
        for agent in self.states:
            state = self.get_state(agent)
            all_agents.append([agent.state.position.x, agent.state.position.y, state.step_count])
        return all_agents
        
    def calc_euclidean(self, loc1, loc2):
        x1 = loc1[0]
        y1 = loc1[1]
        x2 = loc2[0]
        y2 = loc2[1]
        dist = sqrt((x1-x2) ** 2 + (y1-y2) ** 2)
        return dist
        
    def sense_crumbs(self, sensors, num_sensors, start_sensor, agent, cluster_list):
        """
        Generate a (big) array of observations, num_sensors for each crumb
        and store them inside sensors starting at start_sensor.
        Each crumb is stored as: (x,y,exists?,reward)
        """ 
        
        if len(cluster_list) == 0:
            cluster_list = []
            for c in self.crumbs:
                if (c.x, c.y) in getMod().marker_map:
                    cluster_list.append([c.x, c.y])
        i = start_sensor
        closest = None
        closest_distance = None
        pos = agent.state.position
        pos = (pos.x, pos.y)
        for pellet in cluster_list:
            sensors[i] = pellet[0]
            sensors[i+1] = pellet[1]
            if (pellet[0], pellet[1]) in getMod().marker_map:
                sensors[i+2] = 1
                distance = sqrt( (pellet[0] - pos[0]) ** 2 + (pellet[1] - pos[1]) ** 2 )
                if closest is None or distance < closest_distance:
                    closest = pellet
                    closest_distance = distance
            else:
                sensors[i+2] = 0
                cluster_list.remove(pellet)
            sensors[i+3] = 1#pellet.reward
            i = i + num_sensors
        if closest is not None:
            # freebie for scripted agents: tell agent the closest crumb!
            sensors[3] = closest[0]
            sensors[4] = closest[1]
        return True
                     
    def is_active(self, agent):
        """ return true when the agent should act """
        return True
    
    def is_episode_over(self, agent):
        """ is the current episode over for the agent? """
        # Episode ends after number of timesteps
        state = self.get_state(agent)
        if self.EMPTY_ROOM_CONDITION or state.step_count >= self.max_steps:
            num_crumbs = 0
            for c in self.crumbs:
                if (c.x, c.y) in getMod().marker_map:
                    num_crumbs += 1
            return num_crumbs == 0
        else:
            agents = self.get_all_agents()
            for agent in agents:
                if agent[2] >= self.max_steps:
                    return True
            return False
        #if self.max_steps != 0 and state.step_count >= self.max_steps:
        #    return True
        #return False
    
    def cleanup(self):
        """
        cleanup the world
        """
        self.environment = None
        return True

gMod = None

def delMod():
    global gMod
    gMod = None

def getMod():
    global gMod
    if not gMod:
        gMod = SandboxMod()
    return gMod

def ServerMain():
    print "Starting Sandbox Environment"